
# -----------------------------------------------------------------------------
# USEFUL FUNCTIONS for QTL workshop - Uppsala 2011 5-7 April
#
# Author: dario beraldi <dario.beraldi@ed.ac.uk>
# -----------------------------------------------------------------------------

## Use
## source('U:/Documents/uppsala_2011/scripts/qtl_functions.R')
## source('/nfs_netapp/dberaldi/Documents/uppsala_2011/scripts/qtl_functions.R')
## to load them
##
## Note: many functions are not used during the practical, they have been used
## by DB during the preparation of the tutorial itself.


## Parse QMSim output to fit pedigree viewer

split_line<- function(line){
    ## Read a line with fields separated by one or more spaces and return a vector
    line<- strsplit(line, ' ')[[1]]
    line<- line[line != ""]
    return(line)
}
## Read a 'data' file from QMSim to dataframe (e.g. p1_data_001.txt)
read.qsim.data<- function(qsim_file){
    ped_file<- readLines(qsim_file)[-(1:3)]
    header<- split_line(ped_file)
    ped_file<- lapply(ped_file[-1], split_line)
    ped_file<- as.data.frame(do.call("rbind", ped_file), stringsAsFactors= FALSE)
    for(i in c(1:4, 6:13)){
        ped_file[,i]<- as.numeric(ped_file[,i])
    }
    colnames(ped_file)<- header
    return(ped_file)
}

## Read marker file from QMSim (e.g. p1_mrk_001.txt)
read.qsim.mrk<- function(qsim_file){
    mrk_file<- readLines(qsim_file)[-(1:3)]
    header_pdg<- split_line(mrk_file)[1:3]
    mrk_file<- lapply(mrk_file[-1], split_line)
    mrk_file<- as.data.frame(do.call("rbind", mrk_file), stringsAsFactors= FALSE)
    mrk_file<- apply(mrk_file, 2, as.numeric)
    nmrk<- (ncol(mrk_file)-3)/2
    markers<- paste('M', rep(1:((nmrk)), each= 2), rep(c('.1', '.2'), nmrk), sep= '')
    colnames(mrk_file)<- c(header_pdg, markers)
    return(mrk_file)
}

## Expand marker file to have un-genotyped progenies
## ped and mrk are the output of read.qsim.data and read.qsim.mrk repsectively
expand_mrk<- function(ped, mrk){
    gen_file<- merge(ped_file[, c("Progeny", "Sire", "Dam", "Sex")],
                     mrk_file[,-c(2,3)],
                     by= intersect("Progeny", "Progeny"), all.x= TRUE )
    gen_file[is.na(gen_file)]<- 0
    return(gen_file)
}

read.qmsim.map<- function(lm_mrk){
    ## Read the output of QMSim correpsonding to the marker positions
    qmsim<- readLines(lm_mrk)
    qmsim<- qmsim[7:length(qmsim)]
    qmsim_list<- lapply(qmsim, split_line)
    qmsim_df<- as.data.frame(do.call("rbind", qmsim_list), stringsAsFactors= FALSE)
    colnames(qmsim_df)<- c('marker', 'chr', 'position') 
    qmsim_df$position<- as.numeric(qmsim_df$position)
    return(qmsim_df)
}

crimap_format<- function(gen_data, outfile){
    ## Format expanded marker file for crimap
    ## gen_data: Genotype file output of expand_mrk: ID, Sire, Dam, Sex (F/M or 0/1), genotypes
    ## outfile: output file

    ## n_pdg: Number of families- always 1 pedigree.
    n_pdg<- 1
    write.table(n_pdg, outfile, row.names= FALSE, col.names= FALSE, sep= '\t', quote= FALSE)
    nmrk<- (ncol(gen_data)-4)/2
    write.table(nmrk, outfile, row.names= FALSE, col.names= FALSE, sep= '\t', quote= FALSE, append= TRUE)
    
    markers<- sub('.1', '', colnames(gen_data)[-c(1:4)], fixed= TRUE)
    markers[seq(2, length(markers), by= 2)]<- ''
    marker_line<- c(rep('', 4), markers)
    write.table(t(as.matrix(marker_line)), outfile, row.names= FALSE, col.names= FALSE, sep= '\t', quote= FALSE, append= TRUE)
    
    family_id<- 'Pedigree'
    write.table(family_id, outfile, row.names= FALSE, col.names= FALSE, sep= '\t', quote= FALSE, append= TRUE)
    
    pedigree_size<- nrow(gen_data)
    write.table(pedigree_size, outfile, row.names= FALSE, col.names= FALSE, sep= '\t', quote= FALSE, append= TRUE)
    ## Blank line
    write.table('', outfile, row.names= FALSE, col.names= FALSE, sep= '\t', quote= FALSE, append= TRUE)
    
    gen_data[,4][gen_data[,4] == 'F']<- 0
    gen_data[,4][gen_data[,4] == 'M']<- 1
    gen_data[,4][gen_data[,4] != '0' & gen_data[,4] != '1']<- 3
    ## Swap Sire and Dam for crimap
    gen_data<- gen_data[, c(1,3,2, 4:ncol(gen_data))]
    write.table(gen_data, outfile, row.names= FALSE, col.names= FALSE, sep= '\t', quote= FALSE, append= TRUE)    
}


read.crimap.fixed<- function(fixed_out){
    ## Read the output 'crimap fixed' to extract the marker distances
     ## Note: fixed run as Sex_averaged
    ## Output file of crimap fixed e.g.
    ## ./crimap 1 fixed > chr1.fixed
    crimap.fixed<- readLines(fixed_out)
    map_start<- which(crimap.fixed == "Sex_averaged map (recomb. frac., Kosambi cM):")+2
    map_end<- which(crimap.fixed == "* denotes recomb. frac. held fixed in this analysis")-2
    crimap.fixed<- crimap.fixed[map_start:map_end]
    crimap.fixed<- crimap.fixed[seq(1, length(crimap.fixed), by= 2)]
    map<- lapply(crimap.fixed, split_line)
    map<- as.data.frame(do.call("rbind", map), stringsAsFactors= FALSE)
    colnames(map)<- c('mrk_order', 'marker', 'position')
    map$mrk_order<- as.numeric(map$mrk_order)
    map$position<- as.numeric(map$position)
    return(map)
}

plot.crimap.fixed<- function(crimap.fixed, ...){
    ## Plot the map as produced by read.crimap.fixed
    mrk_dist<- crimap.fixed$position
    plot(1,1, ylim= c(min(mrk_dist), max(mrk_dist)), type= 'n', xaxt= 'n', yaxt= 'n', bty= 'n', xlab= '', ylab= '', ...)
    segments(x0=1, x1= 1, y0= min(mrk_dist), y1= max(mrk_dist), lwd= 3, col= 'dodgerblue')
    segments(x0=0.975, x1= 1.025, y0= mrk_dist, y1= mrk_dist, lwd= 3, col= 'dodgerblue')
    text(x= 0.95, y= mrk_dist, crimap.fixed$marker, adj= 1, col= 'grey25')
    text(x= 1.05, y= mrk_dist, crimap.fixed$position, adj= 0, col= 'grey25')
    text(x= 1.125, y= mrk_dist[1:(length(mrk_dist)-1)] + diff(crimap.fixed$position)/2, diff(crimap.fixed$position), adj= 0, col= 'firebrick4')
}

read.crimap.twopoint<- function(crimap.twopoint){
    ## Read the output of crimap twopoint to dataframe. This works only for sex-averaged map.
    twopoint_raw<- readLines(crimap.twopoint)
    twopoint_lods<- twopoint_raw[min(grep('rec. fracs.=', twopoint_raw)):length(twopoint_raw)]
    header<- lapply(twopoint_lods[seq(1, length(twopoint_lods), by= 3)], split_line)
    twopoint_df<- do.call('rbind', header)[,c(1,2,5,8)]
    lods<- lapply(twopoint_lods[seq(2, length(twopoint_lods), by= 3)], split_line)
    lods_df<- do.call('rbind', lods)
    twopoint_df<- as.data.frame(cbind(twopoint_df, lods_df), stringsAsFactors = FALSE)
    colnames(twopoint_df)<- c('marker_map', 'marker_inserted', 'best_r', 'lods', paste('r', c(0.001, 0.01, 0.05, seq(0.1, 0.5, by= 0.05)), sep=''))
    for(i in 4:ncol(twopoint_df)){
        twopoint_df[,i]<- as.numeric(twopoint_df[,i])
    }
    return(twopoint_df)
}

## twopoint<- read.crimap.twopoint('chr1.twopoint')
plot.twopoint<- function(twopoint, positions= NULL, ...){
    ## Plot the LOD score produced by crimap twopoint. Input is the dataframe
    ## produced by read.crimap.twopoint (and only one ,arker inserted)
    ## positions is vector of marker positions (e.g. produced by read.crimap.fixed).
    ## If NULL equal distances are plotted
    par(fg= 'grey40', las= 1)
    if (is.null(positions) == TRUE){
        positions<- 0:(nrow(twopoint)-1)*10
    }
    plot(x= positions, y= twopoint$lods, xlab= 'Markers', ylab= 'LOD score',
         main= paste('LOD score vs.', unique(twopoint$marker_inserted)), type= 'h', lwd= 5, col= 'grey40',
         cex= 0.85, xaxt= 'n', ylim= c(0, max(twopoint$lods)), bty= 'l')
    mtext(text= twopoint$marker_map, at= positions, side= 1, line= 0.5, col= 'grey40')
    abline(h= c(max(twopoint$lods), max(twopoint$lods)-2), lty= 'dotted', col= c('firebrick4', 'grey40'))
    ## i<- 1
    
    toplod<- order(twopoint$lods, decreasing=TRUE)
    for(i in toplod[1:2]){
        lods<- as.numeric(twopoint[i,5:15])
        pos<- positions[i]
        twopoint_grid<- sapply(c(0.001, 0.01, 0.05, seq(0.1, 0.45, by= 0.05)), kosambi_cM)
        points(x= twopoint_grid + pos , y= lods, type= 'o', cex= 0.6, col= 'grey50', lty= 'dotted')
        points(x= pos - twopoint_grid, y= lods, type= 'o', cex= 0.6, col= 'grey50', lty= 'dotted')
    }
}

kosambi_cM<- function(recf){
    ## Convert recombination fraction to Kosambi cM:
    ## d = 1/4 ln[(1+2r)/(1-2r)]
    d<- 0.25 * (log( (1 + 2*recf) / ( 1 - 2*recf) ))*100
    return(d)
}   

get.ibd.pos<- function(loki.ibd.file= NULL){
    if(is.null(loki.ibd.file) == TRUE){
        cat("
            DESCRIPTION
            Reads an ibd matrix produced by loki and returns a vector of
            the positions where IBDs have been estimated.")
    }
    ibd<- read.table(loki.ibd.file, fill= TRUE, sep= ' ', row.names= NULL, stringsAsFactors= FALSE)[, 1:3]
    colnames(ibd)<- c('id1', 'id2', 'ibd')
    ## ibd[1:10,]
    positions<- grep('Position', ibd$id1)
    return(as.numeric(ibd[positions, 3]))
}

parse.loki.ibd<- function(loki.ibd= NULL, outdir= getwd(), getpos= NULL ){
if(is.null(loki.ibd) == TRUE){
cat("
DESCRIPTION
Read an IBD matrix ('loki.ibd') produced by loki and outputs .grm files suitable for
ASREml. Each position in loki.ibd is output as a prefix.position.grm
Returns a vector of filenames that have been written out.
E.g. ibd_pos0001.grm; idb_pos0020.grm; etc...

ARGUMENTS
loki.ibd <string>
    ibd file to parse.
outdir <string>
    Output directory
positions <vector>
    Vector of chromosome positions to be extracted. NULL extract every position.
    Use get.ibd.pos() to see which positions have been computed.
")}
    ## NOTE: Individual IDs must be numeric!!
    ibd<- read.table(loki.ibd, fill= TRUE, sep= ' ', row.names= NULL, stringsAsFactors= FALSE)[, 1:3]
    loki.ibd<- basename(loki.ibd)
    colnames(ibd)<- c('id1', 'id2', 'ibd')
    ## ibd[1:10,]
    ibd<- rbind(ibd, c('**Position', '=', 'EOF'))
    positions<- grep('Position', ibd$id1)
    if(is.null(getpos) == TRUE ){
        getpos<- ibd[positions, 3]
    }
    ## i<- 1
    grm_files<- vector()
    for( i in 1 : (length(positions)-1) ){
        position<- as.numeric(ibd[positions[i], 'ibd'])
        if( (position %in% getpos) == FALSE ){
            next
        }
        ibd.pos<- ibd[ (positions[i]+1) : (positions[i+1]-1), ]
        ibd.pos<- as.data.frame(apply(ibd.pos, 2, as.numeric), , stringsAsFactors= FALSE)
        ibd.pos<- ibd.pos[order(ibd.pos$id1, ibd.pos$id2), ]
        npos<- paste(paste(rep('0', 4 - nchar(position)), collapse= ''), position, collapse= '', sep= '')
        prefix<- strsplit(loki.ibd, '.', fixed = TRUE)[[1]]
        prefix<- paste(prefix[ 1 : (length(prefix)-1)], collapse= '.')
        filename<- paste(prefix, npos, '.grm', sep= '')
        filename_path<- file.path(outdir, filename)
        ## filename<- paste(outdir, '/', prefix, npos, '.grm', sep= '')
        write.table(ibd.pos, file= filename_path, sep= '\t', row.names= FALSE, col.names= FALSE, quote= FALSE)
        grm_files[i]<- filename
    }
    return(as.vector(na.omit(grm_files)))
}

make.qtl.as<- function(noqtl.file= NULL, pedigree.file= '.ped', outdir= getwd(), grm.dir= NULL){
    ## Prepare .as files for testing IBD matrices. A vector of .as filenames is returned.

    ## noqtl.file:    The .as file with only polygenic effect (H0)
    ## pedigree.file: The pedigree file (used to identify where to put the .grm file)
    ## outdir:        Output dir for .as files
    ## grm:           Directory containing .grm files from loki output (parse.loki.ibd)
    
    noqtl<- readLines(noqtl.file)
    ## Which column contains the animal ID:
    id<- split_line(noqtl[grep("!P", noqtl)])[1]
    
    ## Edit the model to add the G matrix (search for the !r ):
    giv<- paste('giv(', id, ',1)', sep= '', collapse='')
    model.indx<- grep("!r", noqtl)
    model.qtl<- paste(noqtl[model.indx], giv, collapse= '')
    
    ## Which line has the pedigree file (.grm file will go after it):
    pedline<- grep(pedigree.file, noqtl)
    
    grm.files<- list.files(grm.dir, full.names= TRUE)    
    grm.files<- grm.files[ grep('.grm', grm.files, fixed= TRUE) ]
    
    ## vector where to put the .as filenames
    asfiles<- vector()
    for(grm in grm.files){
        ## Loop thorough grm files to prepare the .as scripts:
        qtl.as<- noqtl
        qtl.as[model.indx]<- model.qtl
        ibd.mat<- paste(grm, '!NSD', collapse= '')
        qtl.as<- append(qtl.as, ibd.mat, after= pedline)
        ## Files will be named the same as the .grm but with .grm replaced by .as
        filename<- basename(grm)
        filename<- paste(substring(filename, 1, nchar(filename)-4), '.as', sep= '', collaspe= '')
        filename.path<- file.path(outdir, filename)
        writeLines(qtl.as, filename.path)
        asfiles<- append(asfiles, filename.path)
    }
    return(asfiles)
}

scan.qtl<- function(as.files= NULL, asreml.path= NULL, ...){
    ## Run ASREml with all the .as files in as.files.
    ## asreml.as: Path to ASREml (if not on path already)
    asreml<- 'asreml -N'
    if(is.null(asreml.path) == FALSE | asreml.path != ''){
        asreml<- paste(asreml.path, asreml, sep='')
    }
    for (asf in as.files){
        system(paste(asreml, asf, collapse=''), ...)
    }
}

loglik.asr<- function(asr.files){
    ## Get the last LogL= from the asr.files in vector asr.files
    ## Scans the file and get the line with highest index matching 'LogL='
    asloglik<- vector()
    for(f in asr.files){
        asr<- readLines(f)
        asr<- sapply(asr, function(x) sub('LogL=\\s', 'LogL=', x, perl= TRUE))
        logline<- split_line( asr[max(grep(' LogL=', asr))] )
        logcol<- logline[ grep('LogL=', logline) ]
        mlog<- as.numeric(sub('LogL=', '', logcol))
        asloglik<- append(asloglik, mlog)
    }
    return(asloglik)
}


varcomp.asr<- function(file.asr){
    ## Extracts the variance components table from the .asr file
    asr<- readLines(file.asr)
    header<- grep(" Source                Model  terms     Gamma     Component    Comp/SE   % C", asr)
    vcend<- min(grep(' Variance', asr[header : length(asr)])) + (header-1)
    ## vcend<- min(which( asr[header : length(asr)] == "")) + header - 2
    vc.list<- lapply(asr[header:vcend], split_line)
    vc.df<- as.data.frame(do.call('rbind', vc.list), stringsAsFactors= FALSE)
    colnames(vc.df)<- vc.df[1,]
    rownames(vc.df)<- vc.df[,1]
    vc.df<- vc.df[-1,-1]
    vc.df$Model<- as.numeric(vc.df$Model)
    vc.df$terms<- as.numeric(vc.df$terms)
    vc.df$Gamma<- as.numeric(vc.df$Gamma)
    vc.df$Component<- as.numeric(vc.df$Component)
    vc.df$"Comp/SE"<- as.numeric(vc.df$"Comp/SE")
    vc.df$"%"<- as.numeric(vc.df$"%")
    return(vc.df)
}


make_loki_ctrl<- function(chr_map= NULL, genotype_file, ctrl.file= loki.cf){
if(is.null(chr_map) == TRUE){
    cat("
DESCRIPTION
Prepares a control file for LOKI given the marker distances in the dataframe
object chr_map.
File is written to ctrl.file

ARGUMENTS
chr_map <dataframe>
    Dataframe with marker names and positions. Columns must be named 'marker',
    'chr', 'position' (one chromosome only).

genotype_file <dataframe>
    File of genotypes as 'id', 'sire', 'dam', 'sex', 'Marker1A1', 'Marker1A2', 'Marker2A1'...
    (columns names don't matter, order does). Sex is 'M' or 'F'. Missing data is 0

ctrl.file <string>
    String with a name for the output file.
")
    }
    stack<- vector()
    for(i in 1:(length(chr_map$marker)-1)){
        stack[i]<- paste(chr_map$marker[i], '[', 'm', i, 'a, ', 'm', i, 'b], ', sep= '')
    }
    nmrk<- length(chr_map$marker)
    stack[nmrk]<- paste(chr_map$marker[nmrk], '[', 'm', nmrk, 'a, ', 'm', nmrk, 'b]', sep= '')
    stack<- paste(stack, collapse= '')
    
    loki.ctrl<- c(
        paste('FILE [skip=1] "', genotype_file, '", id, sire, dam, sx, ',
              paste(paste('m', rep(1:nrow(chr_map), each= 2), rep(c('a', 'b'), times= nrow(chr_map)), sep= '' ), collapse= ', '), ## m1a, m1b, m2a, m2b ...
             sep= ''
        ),
        'PEDIGREE id,sire,dam',
        "SEX sx 'M', 'F'",
        'MISSING "0"',
        'SET prune_option 0',
        paste('MARKER LOCUS', stack),
        paste('LINK "Chromosome 1"', paste(chr_map$marker, collapse= ', '), collapse= '', sep= ', '),
        'OUTPUT "loki.dat"',
        'LOG "loki.log"'
    )
    writeLines(loki.ctrl, ctrl.file)
    }
 
 
 ## prepare parameter file for loki (.pf)
make_loki_par<- function(niter= 100000, out_start= 10, out_freq= 10,
    output.ibd= 'loki.ibd', ibd_grid= c(0, max(chr_map$position), 10),
    chr_map= NULL, par.file= loki.par){
if (is.null(chr_map) == TRUE){
cat("
DESCRIPTION
Prepares a parameter file for LOKI suitable to estimate IBD probababilities
along the linkage group given in chr_map. IBDs are estimated from position 0 to
the end of the chromosome at frequency (step) given by ibd_grid

ARGUMENTS
niter <int 100000>
    Number of MCMC iterations
out_start <int 10>
    Start outputting statistics from this iteration

out_freq <int 10>
    Store statistics every so many iterations

ibd_grid <vector>
    A vector of three number specifying start, end, step for the computation of
    IBD probabilities. Defualt c(0, max, 10) is to start at position 0 to the end
    of the chromosome at  step of 10 cM.

chr_map <dataframe>
    Dataframe with marker names and positions. Columns must be named 'marker',
    'chr', 'position' (one chromosome only).

output.ibd <string 'loki.ibd'>
    String for the name of the IBD output file.

par.file <string 'loki.par'>
    String for the name of the parameter file to be produced.
")
}
    loki.par<- c(
        'ITERATIONS ', formatC(niter, format= 'd'), '\n',
        'START OUTPUT ', formatC(out_start, format= 'd'), '\n',
        'OUTPUT FREQUENCY ', formatC(out_freq, format= 'd'), '\n',
        'OUTPUT IBD FILE "', output.ibd, '"', '\n',
        paste('START POSITION', chr_map$marker[1], formatC(chr_map$position[1], format= 'd')), '\n',
        paste('POSITION', chr_map$marker[2:length(chr_map$marker)], chr_map$position[2:length(chr_map$position)], '\n'),
        paste('ESTIMATE IBD GRID ', ibd_grid[1], ',', ibd_grid[2], ',', ibd_grid[3]), '\n',
        'SET lm_ratio 0.5'
    )
    lokipar<- paste(loki.par, collapse= '')
    writeLines(lokipar, par.file);
}


make_asreml_template<- function(phen.dat= NULL, phen.file= NULL, pedigree.file= NULL, model= NULL, header_line= NULL, file.as= NULL){
    if(is.null(phen.dat) == TRUE){
        cat("
DESCRIPTION
Prepare a template *.as file for ASREML using the phenotypic data, pedigree and model.
Note: Pedigree and data files must have a header!
        ")
        }
    if (is.null(header_line) == TRUE){
        header_line<- 'ASREML control file - polygenic model'
        }
    
    output<- vector()
    output<- c(output, header_line)
    
    n<- 1
    columns<- colnames(phen.dat)
    for( i in seq(1, ncol(phen.dat)) ){
        if(is.numeric(phen.dat[,i]) == TRUE){
            dtype<- ''
        }else{
            dtype<- '!A'
        }
        if(n == 1) {dtype<- '!P'}
        colname<- paste('   ', columns[i], dtype)
        n<- 0
        output<- c(output, colname)
    }
    if (is.null(pedigree.file) == TRUE)
        { pedigree.file<- '' }
    pedigree_line<- paste(pedigree.file, '!ALPHA !SKIP 1')
    output<- c(output, pedigree_line)

    dataline<- paste(phen.file, '!SKIP 1 !MVREMOVE !MAXIT 100')
    output<- c(output, dataline)
        
    if (is.null(model) == TRUE){ model<- '' }
    output<- c(output, model)
    writeLines(output, file.as)
}

